package org.msi.bayesiannetwork;

import org.msi.bayesiannetwork.xsd.BaseNodeType;
import org.msi.bayesiannetwork.xsd.BayesianNetworkType;
import org.msi.bayesiannetwork.xsd.NodeType;

import javax.xml.bind.JAXBContext;
import javax.xml.bind.JAXBElement;
import javax.xml.bind.JAXBException;
import javax.xml.bind.Unmarshaller;
import java.io.File;
import java.util.HashMap;
import java.util.Map;

/**
 * <p>A factory class for {@link BayesianNetwork}.</p>
 *
 * @author Jakub Narloch
 * @version 1.0
 */
public class BayesianNetworkFactory {

    /**
     * <p>Creates new instance of {@link BayesianNetwork} from specified file.</p>
     *
     * @param file the path to file
     *
     * @return the created instance of {@link BayesianNetwork}
     *
     * @throws BayesianNetworkConfigurationException
     *          if any error occurs during creation of bayesian network
     */
    @SuppressWarnings("unchecked")
    public static BayesianNetwork loadFromFile(String file) {

        Map<String, String> nodeVariable = new HashMap<String, String>();
        BayesianNetwork result;

        try {
            // creates jaxb context
            JAXBContext jaxbContext =
                    JAXBContext.newInstance(org.msi.bayesiannetwork.xsd.ObjectFactory.class);

            // instances marshaller
            Unmarshaller unmarshaller = jaxbContext.createUnmarshaller();

            // unmarshalls the root element
            JAXBElement<BayesianNetworkType> root = (JAXBElement<BayesianNetworkType>)
                    unmarshaller.unmarshal(new File(file));

            // retrieves the root element
            BayesianNetworkType bayesianNetworkDefinition = root.getValue();

            // creates instance of bayesian network
            result = new BayesianNetwork();

            // iterates over all nodes and adds them to network
            if (bayesianNetworkDefinition.getNodes() != null &&
                    bayesianNetworkDefinition.getNodes().getNode() != null) {

                for (NodeType nodeType : bayesianNetworkDefinition.getNodes().getNode()) {

                    nodeVariable.put(nodeType.getName(), nodeType.getVariable());
                }

                for (NodeType nodeType : bayesianNetworkDefinition.getNodes().getNode()) {

                    if (nodeType.getProbability() != null) {
                        // adds simple node
                        result.addNode(nodeType.getName(), nodeType.getVariable(),
                                nodeType.getProbability());
                    } else if (nodeType.getConditionalProbability() != null) {
                        // adds node with conditional probability
                        result.addConditionalNode(nodeType.getName(), nodeType.getVariable(),
                                new ConditionalProbability(Utils.getVariableNamesFromNodes(
                                        nodeType.getConditionalProbability().getNodes().getNode(), nodeVariable),
                                        nodeType.getConditionalProbability().getProbabilities()
                                                .getProbability()));
                    }
                }

                // iterates over all nodes and adds edges to graph
                for (BaseNodeType nodeType : bayesianNetworkDefinition.getNodes().getNode()) {

                    if (nodeType.getEdges() != null && nodeType.getEdges().getTo() != null) {

                        for (String edge : nodeType.getEdges().getTo()) {
                            result.addEdge(nodeType.getName(), edge);
                        }
                    }
                }
            }

            // returns result
            return result;

        } catch (JAXBException e) {

            throw new BayesianNetworkConfigurationException(
                    "Error occurred when creating the bayesian network from configuration.", e);
        } catch (NodeNotFoundException e) {

            throw new BayesianNetworkConfigurationException(
                    "Error occurred when creating the bayesian network from configuration.", e);
        }
    }
}
